<?php
/**
 * Project: AIGC
 * Author: moufer <moufer@163.com>
 * Date: 2024/7/13 13:42
 */

namespace module\ai\library\channel\abstract;

use module\ai\library\channel\interface\CompletionInterface;

abstract class Completion implements CompletionInterface
{
    /**
     * 回答数据
     * @var array
     */
    protected array $data = [
        'message'       => [],
        'stream'        => true,
        'index'         => -1,
        'status'        => 'receiving', //waiting,receiving,error,finish
        'finish_reason' => '',
        'clear_receive' => false,
        'tokens'        => [
            'send_tokens'    => 0,
            'receive_tokens' => 0,
            'total_tokens'   => 0,
        ],
    ];

    /**
     * @param array $rawData 原始数据
     * @param bool $stream 是否流式回答
     * @param string $model 模型名称
     */
    public function __construct(protected array  $rawData,
                                protected bool   $stream = false,
                                protected string $model = '')
    {
        $this->data['stream'] = $this->stream;
        $this->rawToData();
    }

    /**
     * 与另外的$completion合并
     * 用于流式回答时，后续的回答合并到第一个回答中
     * @param Completion $completion
     * @param array $options
     * @return Completion
     */
    public function merge(Completion $completion, array $options = []): static
    {
        $myData = $this->data;
        $data = $completion->data();
        $tmp = array_merge($myData, $data);

        $message = $data['message'];
        if ($data['clear_receive'] !== true) {
            $message['content'] = $myData['message']['content'] . $message['content'];
        }
        $tmp['message'] = $message;

        //出现回答异常时，总token数会变成0，此时就不能覆盖原来的tokens
        if (!$data['tokens']['total_tokens']) {
            $tmp['tokens'] = $myData['tokens'];
        }
//        $tokens = $data['tokens'];
//        foreach ($tokens as $k => $v) {
//            $tokens[$k] = $v + $myData['tokens'][$k];
//        }
//        $tmp['tokens'] = $tokens;

        $this->data = $tmp;
        return $this;
    }

    /**
     * @inheritdoc
     */
    public function raw(): array
    {
        return $this->rawData;
    }

    public function data(): array
    {
        return $this->data;
    }

    public function setCancel(): static
    {
        $this->data['status'] = 'cancel';
        return $this;
    }

    public function setContent($content = ''): static
    {
        $this->data['message']['content'] = $content;
        return $this;
    }

    /**
     * 生成sse数据
     * @param bool $json
     * @return string|array
     */
    public function sse(bool $json = true): string|array
    {
        $data = [
            'status'  => $this->data['status'],
            'content' => $this->data['message']['content'],
            'index'   => $this->data['index'],
            'raw' => $this->raw(),
        ];
        //是否要求前端清空之前返回的信息
        if (!empty($this->data['clear_receive'])) {
            $data['clear'] = true;
        }
        return $json ? json_encode($data, JSON_UNESCAPED_UNICODE) : $data;
    }

    /**
     * @return \Generator
     */
    public function sseChunks(): \Generator
    {
        $sse = $this->sse(false);
        if (empty($sse['content']) || in_array($sse['status'], ['cancel', 'finish'])) {
            yield json_encode($sse, JSON_UNESCAPED_UNICODE);
        } else {
            $contents = mb_str_split($sse['content']);
            $length = count($contents);
            $loop = 0;
            foreach ($contents as $content) {
                $data = array_merge($sse, ['content' => $content]);
                //如果是最后一次循环，就设置结束状态
                if ($data['status'] === 'finish' && $loop < $length - 1) {
                    $data['status'] = 'receiving';
                }
                //$data['cuk_index'] = $loop++;
                yield json_encode($data, JSON_UNESCAPED_UNICODE);
            }
        }
    }

    public function message(): array
    {
        return $this->data['message'] ?? [];
    }

    public function isEnd(): bool
    {
        return !empty($this->data['finish']);
    }

    public function tokens(string $type): int
    {
        return $this->data['tokens'][$type] ?? 0;
    }

    public function errcode(): ?string
    {
        if ($this->data['finish_reason'] === 'error') {
            return $this->data['message']['code'] ?? '0';
        }
        return null;
    }

    protected function error($message, string $code = '0', $clear = false): void
    {
        $this->data['message'] = [
            'content' => $message,
            'role'    => 'error',
            'code'    => $code
        ];
        $this->data['status'] = 'error';
        $this->data['finish_reason'] = 'error';
        //要在前端清除本次回答内容
        if ($clear) {
            $this->data['clear_receive'] = true;
        }
    }

    /**
     * 原始数据转数据
     * @return void
     */
    abstract protected function rawToData(): void;

    /**
     * 设置原始数据
     * @param array $data
     * @param callable|null $callback
     * @return $this
     */
    public function withRaw(array $data, callable $callback = null): static
    {
        $this->rawData = $data;
        if ($callback) {
            $callback($this);
        }
        return $this;
    }

    /**
     * 回答内容
     * @param string $content
     * @param string $role
     * @return $this
     */
    public function withMessage(string $content, string $role): static
    {
        $this->data['message'] = [
            'content' => $content,
            'role'    => $role,
        ];
        return $this;
    }

    /**
     * 流式回答序号
     * @param int $index
     * @return $this
     */
    public function withIndex(int $index): static
    {
        $this->data['index'] = $index;
        return $this;
    }

    /**
     * 流式回答是否最后一条信息
     * @param bool $end
     * @return $this
     */
    public function withIsEnd(bool $end): static
    {
        $this->data['end'] = $end;
        return $this;
    }

    /**
     * 回答总token数
     * @param int $totalTokens
     * @return $this
     */
    public function withTokens(int $totalTokens): static
    {
        $this->data['total_tokens'] = $totalTokens;
        return $this;
    }

    public function toArray(): array
    {
        return $this->data;
    }

    public function geRawData(): array
    {
        return $this->rawData;
    }

    public function getMessage(): string
    {
        return $this->data['message']['content'] ?? '';
    }

    public function __toString(): string
    {
        return json_encode($this->data, JSON_UNESCAPED_UNICODE);
    }
}